// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <gtest/gtest.h>
#include <string>
#include <pollux/common/base/fs.h>
#include <pollux/dwio/common/file_sink.h>
#include <pollux/dwio/common/reader.h>
#include <pollux/testing/dwio/data_files.h>
#include <pollux/dwio/parquet/reader/page_reader.h>
#include <pollux/dwio/parquet/reader/parquet_reader.h>
#include <pollux/dwio/parquet/writer/writer.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::parquet {
    class ParquetTestBase : public testing::Test,
                            public pollux::VectorBuilder {
    protected:
        static void SetUpTestCase() {
            memory::MemoryManager::testingSetInstance({});
        }

        void SetUp() override {
            dwio::common::LocalFileSink::registerFactory();
            rootPool_ = memory::memoryManager()->addRootPool("ParquetTests");
            leafPool_ = rootPool_->addLeafChild("ParquetTests");
            tempPath_ = exec::test::TempDirectoryPath::create();
        }

        static RowTypePtr sampleSchema() {
            return ROW({"a", "b"}, {BIGINT(), DOUBLE()});
        }

        static RowTypePtr dateSchema() {
            return ROW({"date"}, {DATE()});
        }

        static RowTypePtr intSchema() {
            return ROW({"int", "bigint"}, {INTEGER(), BIGINT()});
        }

        static RowTypePtr upperSchemaToLowerCase() {
            return ROW({"a", "b"}, {BIGINT(), BIGINT()});
        }

        std::unique_ptr<kumo::pollux::parquet::ParquetReader> createReader(
            const std::string &path,
            const dwio::common::ReaderOptions &opts) {
            auto input = std::make_unique<dwio::common::BufferedInput>(
                std::make_shared<LocalReadFile>(path), opts.memoryPool());
            return std::make_unique<kumo::pollux::parquet::ParquetReader>(
                std::move(input), opts);
        }

        dwio::common::RowReaderOptions getReaderOpts(
            const RowTypePtr &rowType,
            bool fileColumnNamesReadAsLowerCase = false) {
            dwio::common::RowReaderOptions rowReaderOpts;
            rowReaderOpts.select(
                std::make_shared<kumo::pollux::dwio::common::ColumnSelector>(
                    rowType,
                    rowType->names(),
                    nullptr,
                    fileColumnNamesReadAsLowerCase));

            return rowReaderOpts;
        }

        std::shared_ptr<pollux::common::ScanSpec> makeScanSpec(
            const RowTypePtr &rowType) {
            auto scanSpec = std::make_shared<pollux::common::ScanSpec>("");
            scanSpec->addAllChildFields(*rowType);
            return scanSpec;
        }

        using FilterMap =
        std::unordered_map<std::string, std::unique_ptr<pollux::common::Filter> >;

        void assertEqualVectorPart(
            const VectorPtr &expected,
            const VectorPtr &actual,
            vector_size_t offset) {
            ASSERT_GE(expected->size(), actual->size() + offset);
            ASSERT_EQ(expected->type_kind(), actual->type_kind());
            for (vector_size_t i = 0; i < actual->size(); i++) {
                ASSERT_TRUE(expected->equal_value_at(actual.get(), i + offset, i))
          << "at " << (i + offset) << ": expected "
          << expected->toString(i + offset) << ", but got "
          << actual->toString(i);
            }
        }

        void assertReadWithReaderAndExpected(
            std::shared_ptr<const RowType> outputType,
            dwio::common::RowReader &reader,
            RowVectorPtr expected,
            memory::MemoryPool &memoryPool) {
            uint64_t total = 0;
            VectorPtr result = BaseVector::create(outputType, 0, &memoryPool);
            while (total < expected->size()) {
                auto part = reader.next(1000, result);
                if (part > 0) {
                    assertEqualVectorPart(expected, result, total);
                    total += result->size();
                } else {
                    break;
                }
            }
            EXPECT_EQ(total, expected->size());
            EXPECT_EQ(reader.next(1000, result), 0);
        }

        void assertReadWithReaderAndFilters(
            const std::unique_ptr<dwio::common::Reader> reader,
            const std::string & /* fileName */,
            const RowTypePtr &fileSchema,
            FilterMap filters,
            const RowVectorPtr &expected) {
            auto scanSpec = makeScanSpec(fileSchema);
            for (auto &&[column, filter]: filters) {
                scanSpec->getOrCreateChild(pollux::common::Subfield(column))
                        ->setFilter(std::move(filter));
            }

            auto rowReaderOpts = getReaderOpts(fileSchema);
            rowReaderOpts.setScanSpec(scanSpec);
            auto rowReader = reader->createRowReader(rowReaderOpts);
            assertReadWithReaderAndExpected(
                fileSchema, *rowReader, expected, *leafPool_);
        }

        std::unique_ptr<dwio::common::FileSink> createSink(
            const std::string &filePath) {
            auto sink = dwio::common::FileSink::create(
                fmt::format("file:{}", filePath), {.pool = rootPool_.get()});
            EXPECT_TRUE(sink->isBuffered());
            EXPECT_TRUE(fs::exists(filePath));
            EXPECT_FALSE(sink->isClosed());
            return sink;
        }

        std::unique_ptr<kumo::pollux::parquet::Writer> createWriter(
            std::unique_ptr<dwio::common::FileSink> sink,
            std::function<
                std::unique_ptr<kumo::pollux::parquet::DefaultFlushPolicy>()>
            flushPolicy,
            const RowTypePtr &rowType,
            kumo::pollux::common::CompressionKind compressionKind =
                    kumo::pollux::common::CompressionKind_NONE) {
            kumo::pollux::parquet::WriterOptions options;
            options.memoryPool = rootPool_.get();
            options.flushPolicyFactory = flushPolicy;
            options.compressionKind = compressionKind;
            return std::make_unique<kumo::pollux::parquet::Writer>(
                std::move(sink), options, rowType);
        }

        std::vector<RowVectorPtr> createBatches(
            const RowTypePtr &rowType,
            uint64_t numBatches,
            uint64_t vectorSize) {
            std::vector<RowVectorPtr> batches;
            batches.reserve(numBatches);
            VectorFuzzer fuzzer({.vectorSize = vectorSize}, leafPool_.get());
            for (auto i = 0; i < numBatches; ++i) {
                batches.emplace_back(fuzzer.fuzzInputFlatRow(rowType));
            }
            return batches;
        }

        std::string getExampleFilePath(const std::string &fileName) {
            return test::getDataFilePath(
                "pollux/dwio/parquet/tests/reader", "../examples/" + fileName);
        }

        static constexpr uint64_t kRowsInRowGroup = 10'000;
        static constexpr uint64_t kBytesInRowGroup = 128 * 1'024 * 1'024;
        std::shared_ptr<memory::MemoryPool> rootPool_;
        std::shared_ptr<memory::MemoryPool> leafPool_;
        std::shared_ptr<exec::test::TempDirectoryPath> tempPath_;
    };
} // namespace kumo::pollux::parquet
